from lib.StockDataFetcher import StockDataFetcher
import pandas as pd
import matplotlib.pyplot as plt
import mpl_finance as mpf
import matplotlib.dates as mdates

# 使用示例
if __name__ == "__main__":
    fetcher = StockDataFetcher('000737',start_date='2024-07-01',forced_update=True)
    # print(df.columns) # trade_time', 'open', 'close', 'volume', 'high', 'low', 'amount', 'change', 'change_pct', 'turnover_ratio', 'pre_close', 'stock_code','trade_date'
    df = fetcher.fetch_or_load_data()[['trade_date', 'open', 'close', 'high', 'low']]
    df = df.rename(columns={'trade_date': 'Date', 'open': 'Open', 'close': 'Close', 'high': 'High', 'low': 'Low','volume':'Volume'})
    print(df)

    df['Date'] = pd.to_datetime(df['Date'])  # 转换日期列到datetime类型

    df['Open'] = df['Open'].astype(float)
    df['Close'] = df['Close'].astype(float)
    df['High'] = df['High'].astype(float)
    df['Low'] = df['Low'].astype(float)

  

    # 准备数据以供mpl_finance使用
    ohlc = df.loc[:, ['Date', 'Open', 'High', 'Low', 'Close']].values
    ohlc[:, 0] = mdates.date2num(ohlc[:, 0])  # 将日期转换为matplotlib可以理解的数字

    # 绘制K线图
    fig, ax = plt.subplots()
    mpf.candlestick_ochl(ax, ohlc,  colorup='g', colordown='r')

    # 设置x轴的日期格式
    ax.xaxis_date()
    plt.xticks(rotation=45)

    # 显示图形
    plt.show()